# 用于SSFM光纤每段仿真后的EDFA放大,并加入噪声
#   sig：输入需要放大的信号，复值形式,列向量或者数组
#   h:普朗克常数。
#   polar_n:偏振数,用于噪声功率的分配,双偏振情况下功率减半。
#   nf_db:噪声系数，dBm单位
#   B:信号带宽，Hz单位
#   f_carrier:被放大信号的载波频率,Hz
#   alpha_db:衰减系数，5dB/km单位
#   span_len:span长度km
import numpy as np


def Amplifier(sig, Amplifier_para, Signal, Fiber_para):
    h = 6.626068e-34
    polar_n = Signal['polarization']
    nf_db = Amplifier_para['EDFA_nf_db']

    # db to unit-less
    nf = 10 ** (nf_db / 10)

    B = Signal['Sampling_rate']
    f_carrier = Signal['f_c']
    alpha_db = Fiber_para['alpha_dB']
    span_len = Fiber_para['L']
    gain_db = alpha_db * span_len

    # db to unit-less
    gain = 10 ** (gain_db / 10)

    data = sig

    # 信号放大
    # unit-less to Np
    g = np.log(gain)

    # 计算噪声功率
    n_pdf = h * f_carrier * (gain - 1) * nf
    n_power = n_pdf * B

    # 对单偏振或者双偏振添加噪声
    if polar_n == 1:

        data = data * np.exp(g / 2)

        # Noise
        noise_ASE_real = np.random.randn(len(data))
        noise_ASE_imag = np.random.randn(len(data))
        n = np.sqrt(n_power) * (noise_ASE_real + 1j * noise_ASE_imag)

        signal = data + n

    return signal
